#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
对 Geant4 mesh 1D 输出 csv 做后处理:
  计算平均值, 方差并画图
  文件格式: iX,iY,iZ,value, value^2, entry
  bin=1x1x150, 因此只按 iZ 分组
"""

import glob, re, os
import numpy as np
import matplotlib.pyplot as plt
import pathlib as pl
from icecream import ic

_CWD = pl.Path(__file__).parent
_POST_DIR = _CWD / "post"
# ic(_POST_DIR)


# ---------- 可调参数 ----------
FIGSIZE = (8, 5)  # 图幅
SAVE_DPI = 200  # 保存分辨率
XLABEL = "z slice (iZ)"
YLABEL_DICT = {
    "eDep": "Energy deposit  / MeV",
    "LETd": "LETd  / (keV/μm)",  # 如果单位不同请自己改
    "rgProj": "Range proj  / mm",  # 同上
}
# -------------------------------


def prepare():
    os.makedirs(_POST_DIR, exist_ok=True)


def parse_csv(path):
    """返回 dict: {iz: (sum_val, sum_val2, n)}"""
    data = {}
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            ix, iy, iz, sv, sv2, ent = line.split(",")
            iz = int(iz)
            data[iz] = (float(sv), float(sv2), int(ent))
    return data


def calc_mean_var(dat):
    """
    输入: {iz: (sum_val, sum_val2, n)}
    输出: (iz_arr, mean_arr, var_arr, sigma_arr)
    """
    izs = sorted(dat.keys())
    means, vars_ = [], []
    for iz in izs:
        sv, sv2, n = dat[iz]
        if n == 0:
            m, v = 0.0, 0.0
        else:
            m = sv / n
            v = (sv2 - sv * sv / n) / n  # 总体方差
        means.append(m)
        vars_.append(v)
    return np.array(izs), np.array(means), np.array(vars_), np.sqrt(vars_)


def plot_one(csv_file, save=True):
    """处理单个文件并画图"""
    basename = pl.Path(csv_file).stem
    # ic(basename)
    # 从文件名猜测 scorer 名字
    scorer = re.findall(r"(\w+)_700\.0MeV", basename)
    scorer = scorer[0] if scorer else "quantity"
    ylabel = YLABEL_DICT.get(scorer, scorer)

    dat = parse_csv(csv_file)
    iz, mean, var, sigma = calc_mean_var(dat)

    plt.figure(figsize=FIGSIZE)
    plt.fill_between(
        iz, mean - sigma, mean + sigma, color="C0", alpha=0.25, label="±1σ"
    )
    plt.plot(iz, mean, lw=1.2, color="C0", label="Mean")
    plt.xlabel(XLABEL)
    plt.ylabel(ylabel)
    plt.title(f"{scorer} @ 700 MeV")
    plt.tight_layout()
    plt.legend()
    if save:
        outpng = (_POST_DIR / basename).with_suffix(".png")
        # ic(outpng)
        plt.savefig(outpng, dpi=SAVE_DPI)
        print(f"Saved  -> {outpng}")
    plt.close()


def main():
    prepare()

    files = glob.glob("out/*700.0MeV.csv")
    if not files:
        print("No *700.0MeV.csv found!")
        return
    for f in files:
        print(f"Processing {f} ...")
        plot_one(f, save=True)


if __name__ == "__main__":
    main()
